/***************************************************************************************************
 * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 *modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright notice,
 *this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *notice, this list of conditions and the following disclaimer in the
 *documentation and/or other materials provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the names of its
 *contributors may be used to endorse or promote products derived from this
 *software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
 *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/*! \file
    \brief Implicit GEMM testbed sizes for Conv2d problem
*/
#pragma once

#include "../../common/cutlass_unit_test.h"

#include "cutlass/cutlass.h"

#include "cutlass/aligned_buffer.h"
#include "cutlass/numeric_types.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/layout/tensor.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/core_io.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/conv/convolution.h"
#include "cutlass/conv/conv2d_problem_size.h"
#include "cutlass/conv/conv3d_problem_size.h"

namespace test {
namespace conv {
namespace device {

using Conv3dProblemVector = std::vector<cutlass::conv::Conv3dProblemSize>;

////////////////////////////////////////////////////////////////////////////
/// Structure TestbedConv3dProblemSizes initializes and holds conv default and
/// important network sizes
////////////////////////////////////////////////////////////////////////////
struct TestbedConv3dProblemSizes {
    //
    // Data members
    //
    int minimum_channel_size;
    Conv3dProblemVector conv3d_default_sizes;
    Conv3dProblemVector conv3d_vnet_medical_sizes;

    //
    // Methods
    //
    /// Default ctor
    TestbedConv3dProblemSizes(int minimum_channel_size_ = 64)
            : minimum_channel_size(minimum_channel_size_) {
        initialize_conv3d_default_sizes();
        initialize_conv3d_vnet_medical_sizes(conv3d_vnet_medical_sizes,
                                             1 /*batch-size*/);

        filter_all();
    }

    /// Eliminates some illegal cases
    void filter_all() {
        Conv3dProblemVector* problems_vectors[] = {&conv3d_default_sizes,
                                                   &conv3d_vnet_medical_sizes};

        for (Conv3dProblemVector* problems : problems_vectors) {
            Conv3dProblemVector filtered;

            for (cutlass::conv::Conv3dProblemSize const& problem : *problems) {
                if (!(problem.C % minimum_channel_size)) {
                    filtered.push_back(problem);
                }
            }

            *problems = filtered;
        }
    }

    // Add a few standard convolution problem sizes
    void initialize_conv3d_default_sizes() {
        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 1, 3, 3, minimum_channel_size},  // input size  (NDHWC)
                {8, 1, 1, 1, minimum_channel_size},  // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 1, 16, 16, minimum_channel_size},  // input size  (NDHWC)
                {8, 1, 3, 3, minimum_channel_size},    // filter size (KTRSC)
                cutlass::Coord<3>({0, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 1, 15, 19, 160},           // input size  (NDHWC)
                {224, 1, 3, 6, 160},           // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 2, 1, 1, minimum_channel_size},  // input size  (NDHWC)
                {8, 2, 1, 1, minimum_channel_size},  // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 1, 7, 7, minimum_channel_size},   // input size  (NDHWC)
                {16, 1, 3, 3, minimum_channel_size},  // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_default_sizes.push_back(cutlass::conv::Conv3dProblemSize(
                {1, 11, 15, 19, 64},           // input size  (NDHWC)
                {32, 4, 3, 6, 64},             // filter size (KTRSC)
                cutlass::Coord<3>({2, 1, 3}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));
    }

    // Add vnet layers to unit testing sizes
    void initialize_conv3d_vnet_medical_sizes(
            Conv3dProblemVector& conv3d_problem_vector, int batch_size = 1) {
        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 32, 32, 32, 16},  // input size  (NDHWC)
                {32, 2, 2, 2, 16},             // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {2, 2, 2}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 16, 16, 16, 32},  // input size  (NDHWC)
                {32, 3, 3, 3, 32},             // filter size (KTRSC)
                cutlass::Coord<3>({1, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 16, 16, 16, 32},  // input size  (NDHWC)
                {64, 2, 2, 2, 32},             // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {2, 2, 2}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 8, 8, 8, 64},     // input size  (NDHWC)
                {64, 3, 3, 3, 64},             // filter size (KTRSC)
                cutlass::Coord<3>({1, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 8, 8, 8, 64},     // input size  (NDHWC)
                {128, 2, 2, 2, 64},            // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {2, 2, 2}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 4, 4, 4, 128},    // input size  (NDHWC)
                {128, 3, 3, 3, 128},           // filter size (KTRSC)
                cutlass::Coord<3>({1, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 8, 8, 8, 128},    // input size  (NDHWC)
                {128, 3, 3, 3, 128},           // filter size (KTRSC)
                cutlass::Coord<3>({1, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 16, 16, 16, 64},  // input size  (NDHWC)
                {64, 3, 3, 3, 64},             // filter size (KTRSC)
                cutlass::Coord<3>({1, 1, 1}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {1, 1, 1}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 32, 32, 32, 16},  // input size  (NDHWC)
                {64, 2, 2, 2, 16},             // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {2, 2, 2}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));

        conv3d_problem_vector.push_back(cutlass::conv::Conv3dProblemSize(
                {batch_size, 16, 16, 16, 32},  // input size  (NDHWC)
                {128, 2, 2, 2, 32},            // filter size (KTRSC)
                cutlass::Coord<3>({0, 0, 0}),  // padding (pad_d, pad_h, pad_w)
                cutlass::Coord<3>(
                        {2, 2, 2}),  // stride (stride_d, stride_h, stride_w)
                cutlass::Coord<3>(
                        {1, 1,
                         1})  // dilation (dilation_d, dilation_h, dilation_w)
                ));
    }
};

}  // namespace device
}  // namespace conv
}  // namespace test
